# 对应AI衣语模块

# -*- coding: utf-8 -*-
import mysql.connector
import pandas as pd
from openai import OpenAI
import sys
import io
import re  # 用于解析返回结果

# 设置标准输出的编码为UTF-8
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

print("Python script started")
print("Arguments received:", sys.argv)

# 设置 OpenAI API 客户端
client = OpenAI(
    base_url="https://api.openai-proxy.org/v1",
    api_key="sk-0CCtNIrBc0s3GJz1JI5GSSh9mZh7ehTARtT6kIY18DlYxlfD",
)

def fetch_clothes_data():
    """从数据库中获取所有衣服数据"""
    try:
        conn = mysql.connector.connect(
            host="127.0.0.1",
            user="root",
            password="SYJwww62442@",
            database="stylefit",
            port=3306
        )
        cursor = conn.cursor()
        query = "SELECT clothes_id, description FROM clothes1;"
        cursor.execute(query)
        result = cursor.fetchall()
        column_names = [desc[0] for desc in cursor.description]
        conn.close()
        df = pd.DataFrame(result, columns=column_names)
        return df
    except mysql.connector.Error as e:
        print(f"错误：{e}")
        return None

def analyze_clothes_with_chatgpt(clothes_data, text):
    """利用 ChatGPT API 分析用户请求，返回符合要求的衣服 ID 或用户请求的答案"""
    if clothes_data is None or clothes_data.empty:
        return None

    # 准备所有衣服的描述信息
    clothes_info = []
    for index, row in clothes_data.iterrows():
        clothes_id = row['clothes_id']
        description = row['description']
        clothes_info.append(f"ID: {clothes_id}, 描述: {description}")

    # 将所有衣服信息格式化为一个字符串
    clothes_info_str = "\n".join(clothes_info)

    # 调用 ChatGPT API 进行分析
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": f"你是一个专业的时尚顾问。用户请求如下：{text}。如果用户希望你推荐衣服，请根据以下衣服信息返回符合要求的所有衣服ID，用逗号分隔：\n{clothes_info_str}。如果用户不需要推荐衣服，请直接回答用户的问题。",
                },
                {
                    "role": "user",
                    "content": "请根据我的请求判断是否需要返回衣服ID。如果是，请返回符合要求的所有衣服ID，用逗号分隔。如果不是，请对我请求的问题进行分析返回答案。"
                }
            ],
            model="gemini-1.5-flash",  # 如果是其他兼容模型，比如 deepseek，直接这里改模型名即可，其他都不用动
        )
        answer = chat_completion.choices[0].message.content.strip()

        # 解析返回结果
        if answer:
            # 检查是否是ID列表
            if re.match(r"^\d+(,\s*\d+)*$", answer):
                # 提取衣服ID
                clothes_ids = answer.split(',')
                return {"type": "clothes_ids", "data": clothes_ids}
            else:
                # 返回用户请求的答案
                return {"type": "answer", "data": answer}

        return {"type": "answer", "data": "没有找到符合条件的衣服或答案。"}

    except Exception as e:
        print(f"调用 ChatGPT API 时出错：{e}")
        return {"type": "error", "data": f"API调用出错：{e}"}

def get_matching_clothes_ids_or_answer(text):
    """主函数，接收用户请求，返回符合要求的衣服 ID 或用户请求的答案"""
    # 获取所有衣服数据
    clothes_data = fetch_clothes_data()

    if clothes_data is None:
        return {"type": "error", "data": "无法获取衣服数据"}

    # 分析并获取结果
    result = analyze_clothes_with_chatgpt(clothes_data, text)

    return result

if __name__ == "__main__":
    print("Main block started")
    if len(sys.argv) > 1:
        text = sys.argv[1]
        print("Received text:", text)
        result = get_matching_clothes_ids_or_answer(text)

        if result["type"] == "clothes_ids":
            print("Matching clothes IDs:", result["data"])
            for clothes_id in result["data"]:
                print(clothes_id)
        elif result["type"] == "answer":
            print("Answer:", result["data"])
        else:
            print("Error:", result["data"])
    else:
        print("No arguments provided")
